import yaml
import os

# 数据集信息
data_dir = './datasets'  # 数据集的根目录
nc = 5  # 类别数，根据实际类别数修改
class_names = ['普通士兵', '工兵', '火箭兵', '指挥官']  # 类别名称

# 设置数据集路径
yaml_data = {
    'path': './',  # 数据集根目录
    'train': 'images/train',
    'val': 'images/val',
    'test': 'images/test',
    'nc': len(class_names),  # 类别数
    'names': class_names  # 类别名称
}

# 将数据写入yaml文件
yaml_file_path = os.path.join(data_dir, 'data.yaml')
with open(yaml_file_path, 'w') as yaml_file:
    yaml.dump(yaml_data, yaml_file, default_flow_style=False)

print(f"已成功创建 {yaml_file_path}")
